#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from enum import Enum
import os

class GeneralRarity(Enum):
    """武将稀有度"""
    COMMON = "普通"       # 1星
    UNCOMMON = "精良"     # 2星  
    RARE = "稀有"         # 3星
    EPIC = "史诗"         # 4星
    LEGENDARY = "传说"    # 5星

class GeneralFaction(Enum):
    """武将阵营"""
    WEI = "魏"
    SHU = "蜀" 
    WU = "吴"
    QUN = "群"

class General:
    """武将类"""
    def __init__(self, name: str, faction: GeneralFaction = GeneralFaction.QUN, 
                 rarity: GeneralRarity = GeneralRarity.COMMON, cost: int = 1):
        self.name = name
        self.faction = faction
        self.rarity = rarity
        self.cost = cost
        
    def __str__(self):
        return f"{self.name}({self.faction.value})"

# 基于游戏资料/武将图片目录的武将枚举
AVAILABLE_GENERALS = [
    General("丁奉", GeneralFaction.WU, GeneralRarity.UNCOMMON, 2),
    General("严政", GeneralFaction.WEI, GeneralRarity.COMMON, 1),
    General("伏完", GeneralFaction.WEI, GeneralRarity.COMMON, 1),
    General("关羽", GeneralFaction.SHU, GeneralRarity.LEGENDARY, 5),
    General("冯芳", GeneralFaction.QUN, GeneralRarity.COMMON, 1),
    General("刘禅", GeneralFaction.SHU, GeneralRarity.RARE, 3),
    General("华雄", GeneralFaction.QUN, GeneralRarity.UNCOMMON, 2),
    General("卢植", GeneralFaction.QUN, GeneralRarity.UNCOMMON, 2),
    General("吕蒙", GeneralFaction.WU, GeneralRarity.EPIC, 4),
    General("吴兰", GeneralFaction.SHU, GeneralRarity.COMMON, 1),
    General("夏侯惇", GeneralFaction.WEI, GeneralRarity.EPIC, 4),
    General("张星彩", GeneralFaction.SHU, GeneralRarity.UNCOMMON, 2),
    General("张横", GeneralFaction.QUN, GeneralRarity.COMMON, 1),
    General("徐氏", GeneralFaction.SHU, GeneralRarity.UNCOMMON, 2),
    General("陶升", GeneralFaction.QUN, GeneralRarity.COMMON, 1),
    General("文丑", GeneralFaction.QUN, GeneralRarity.UNCOMMON, 2),
    General("曹休", GeneralFaction.WEI, GeneralRarity.RARE, 3),
    General("朱儁", GeneralFaction.QUN, GeneralRarity.RARE, 3),
    General("李典", GeneralFaction.WEI, GeneralRarity.UNCOMMON, 2),
    General("桥蕤", GeneralFaction.QUN, GeneralRarity.COMMON, 1),
    General("波才", GeneralFaction.QUN, GeneralRarity.COMMON, 1),
    General("王异", GeneralFaction.WEI, GeneralRarity.RARE, 3),
    General("眭固", GeneralFaction.QUN, GeneralRarity.COMMON, 1),
    General("糜竺", GeneralFaction.SHU, GeneralRarity.UNCOMMON, 2),
    General("荀谌", GeneralFaction.WEI, GeneralRarity.UNCOMMON, 2),
    General("董白", GeneralFaction.QUN, GeneralRarity.UNCOMMON, 2),
    General("董贵人", GeneralFaction.QUN, GeneralRarity.COMMON, 1),
    General("蔡夫人", GeneralFaction.QUN, GeneralRarity.COMMON, 1),
    General("邓艾", GeneralFaction.WEI, GeneralRarity.LEGENDARY, 5),
    General("邓茂", GeneralFaction.QUN, GeneralRarity.COMMON, 1),
    General("邹氏", GeneralFaction.QUN, GeneralRarity.COMMON, 1),
    General("陆逊", GeneralFaction.WU, GeneralRarity.LEGENDARY, 5),
    General("陈武", GeneralFaction.WU, GeneralRarity.UNCOMMON, 2),
    General("陶谦", GeneralFaction.QUN, GeneralRarity.RARE, 3),
    General("雷薄", GeneralFaction.QUN, GeneralRarity.COMMON, 1),
    General("颜良", GeneralFaction.QUN, GeneralRarity.UNCOMMON, 2),
    General("马云騄", GeneralFaction.SHU, GeneralRarity.UNCOMMON, 2),
    General("马休", GeneralFaction.SHU, GeneralRarity.COMMON, 1),
    General("严夫人", GeneralFaction.WEI, GeneralRarity.UNCOMMON, 2),  # 新增武将
]

def get_generals_by_faction(faction: GeneralFaction):
    """根据阵营获取武将"""
    return [g for g in AVAILABLE_GENERALS if g.faction == faction]

def get_generals_by_rarity(rarity: GeneralRarity):
    """根据稀有度获取武将"""
    return [g for g in AVAILABLE_GENERALS if g.rarity == rarity]

def get_general_by_name(name: str):
    """根据名称获取武将"""
    for general in AVAILABLE_GENERALS:
        if general.name == name:
            return general
    return None

def get_all_general_names():
    """获取所有武将名称列表"""
    return [g.name for g in AVAILABLE_GENERALS]

def load_generals_from_directory(directory_path: str = "游戏资料/武将图片"):
    """从目录动态加载武将列表"""
    if not os.path.exists(directory_path):
        return get_all_general_names()
    
    general_names = []
    for filename in os.listdir(directory_path):
        if filename.endswith('.png'):
            name = filename.replace('.png', '')
            general_names.append(name)
    
    return sorted(general_names)

# 武将分组显示
GENERALS_BY_FACTION = {
    "魏国": get_generals_by_faction(GeneralFaction.WEI),
    "蜀国": get_generals_by_faction(GeneralFaction.SHU),
    "吴国": get_generals_by_faction(GeneralFaction.WU),
    "群雄": get_generals_by_faction(GeneralFaction.QUN),
}

GENERALS_BY_RARITY = {
    "传说": get_generals_by_rarity(GeneralRarity.LEGENDARY),
    "史诗": get_generals_by_rarity(GeneralRarity.EPIC),
    "稀有": get_generals_by_rarity(GeneralRarity.RARE),
    "精良": get_generals_by_rarity(GeneralRarity.UNCOMMON),
    "普通": get_generals_by_rarity(GeneralRarity.COMMON),
}

if __name__ == "__main__":
    print("三国杀自走棋武将列表:")
    print(f"总计: {len(AVAILABLE_GENERALS)} 个武将")
    
    print("\n按阵营分类:")
    for faction, generals in GENERALS_BY_FACTION.items():
        print(f"{faction}: {len(generals)}个 - {[g.name for g in generals]}")
    
    print("\n按稀有度分类:")
    for rarity, generals in GENERALS_BY_RARITY.items():
        print(f"{rarity}: {len(generals)}个 - {[g.name for g in generals]}")